import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from datetime import datetime, timedelta
import csv
import bisect

pains = ['无','轻微','中等','严重']
stresss = ['低','中','高']

# 数据预处理类（修改版）
class DateFeatureDataset(Dataset):
    def __init__(self, data):
        """
        data: 包含完整日期特征的数据列表
        """
        # 转换日期为特征向量
        self.features = []  # 单独存储特征
        self.labels = []     # 单独存储标签
        self.max_period = 0
        base_date = datetime(2000, 1, 1)
        for item in data:
            current_date = datetime.strptime(item[0], '%Y-%m-%d')
            cycle_day = (current_date - datetime.strptime(item[4], '%Y-%m-%d')).days  # 当前周期第几天
            self.max_period = max(self.max_period, cycle_day)

        print(self.max_period)
        for item in data:
            current_date = datetime.strptime(item[0], '%Y-%m-%d')
            self.features.append((current_date - base_date).days % self.max_period)
            self.labels.append(item[1])

        self.mean = np.mean(self.features)
        self.std = np.std(self.features)

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        x = self.features[idx]
        y = self.labels[idx]
        x = (x - self.mean) / self.std
        return (
            torch.FloatTensor([x]),  # 形状 [1]
            torch.FloatTensor([y])
        )

# 全连接网络模型
class DatePredictor(nn.Module):
    def __init__(self, input_size=1, hidden_size=64):
        super().__init__()
        # 共享特征提取层
        self.base_net = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU()
        )
        # 多任务输出头
        self.menstruation_head = nn.Sequential(
            nn.Linear(hidden_size, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        features = self.base_net(x)
        return self.menstruation_head(features)

# 使用示例
if __name__ == "__main__":
    # 加载数据
    data = []
    with open('period_full_data.csv', 'r', encoding='utf-8') as f:
        reader = csv.reader(f)
        next(reader)  # 跳过标题
        for row in reader:
            data.append([
                row[0],  # 日期
                int(row[1]),  # 是否经期（标签）
                pains.index(row[2]),  # 疼痛程度（标签）
                stresss.index(row[3]),# 压力水平（标签）
                row[4],#上次开始日期
            ])

    print(data)

    # 创建数据集
    dataset = DateFeatureDataset(data)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

    # 初始化模型
    model = DatePredictor()
    criterion_m = nn.BCELoss()  # 二分类
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 训练循环
    for epoch in range(1000):
        for x, y in dataloader:
            optimizer.zero_grad()
            m_pred = model(x)
            loss_m = criterion_m(m_pred, y)
            loss_m.backward()
            optimizer.step()
        if epoch % 100 == 0:
            print(f'Epoch {epoch}, Loss: {loss_m.item():.4f}')

    # 预测任意日期
    def predict_date(target_date):
        base_date = datetime(2000, 1, 1)
        target_dt = datetime.strptime(target_date, '%Y-%m-%d')

        # 获取特征值并归一化
        raw_feature = (target_dt - base_date).days % dataset.max_period
        norm_feature = (raw_feature - dataset.mean) / dataset.std

        # 创建正确维度的张量 [batch_size=1, input_size=1]
        input_tensor = torch.FloatTensor([[norm_feature]])

        with torch.no_grad():
            m_pred = model(input_tensor)

        is_menstruating = int(m_pred.item() > 0.5)

        return {
            'date': target_date,
            'is_menstruating': is_menstruating,
        }

    # 示例预测
    print(predict_date('2023-10-27'))
    print(predict_date('2023-10-28'))
    print(predict_date('2023-10-29'))
    print(predict_date('2023-10-30'))

